import json
import base64
from fastapi import APIRouter, Depends, Header, Response, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse
from langchain_core.messages import HumanMessage, AIMessage
from typing import Optional

from app.models.chat import ChatRequest
from app.services.rag.retrieval import MilvusRetriever
from app.services.tools.agent import ToolAgent
from app.services.tools.agent_manager import get_singleton_agent, get_singleton_llm
from app.services.tools.memory_manager import get_memory_manager
from app.services.vector_store.milvus_service import MilvusService
from app.services.tts.edge_tts_service import tts_service
from app.services.audio_manager import audio_manager

router = APIRouter(
    prefix="/chat",
    tags=["chat"]
)

# 依赖项：获取LLM模型（单例）
def get_llm():
    return get_singleton_llm()

# 依赖项：获取工具代理（单例）
def get_tool_agent():
    return get_singleton_agent()

# 依赖项：获取Milvus服务（用于RAG）
def get_milvus_service():
    from langchain_openai import OpenAIEmbeddings
    return MilvusService(embedding_model=OpenAIEmbeddings())

# 依赖项：获取记忆管理器
def get_memory():
    return get_memory_manager()

@router.websocket("/ws")
async def websocket_chat(
    websocket: WebSocket,
):
    """WebSocket聊天接口，支持流式输出、RAG和工具使用"""
    await websocket.accept()
    
    # 获取依赖项
    tool_agent = get_tool_agent()
    milvus_service = get_milvus_service()
    memory_manager = get_memory()
    
    # 初始化session_id为None，等待客户端发送
    session_id = None
    
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()
            
            # 解析消息类型
            message_type = data.get("type", "chat")
            
            if message_type == "session_info":
                # 处理session_info消息，提取session_id
                session_id = data.get("session_id")
                if session_id:
                    # 确认session_id已接收
                    await websocket.send_json({
                        "type": "session_confirmed",
                        "session_id": session_id,
                        "message": "Session ID已设置"
                    })
                else:
                    # 如果没有提供session_id，生成一个新的
                    session_id = memory_manager.generate_session_id()
                    await websocket.send_json({
                        "type": "session_info",
                        "session_id": session_id,
                        "message": "已生成新的Session ID"
                    })
                    
            elif message_type == "chat":
                # 检查是否已设置session_id
                if not session_id:
                    await websocket.send_json({
                        "type": "error",
                        "message": "请先发送session_info消息设置session_id"
                    })
                    continue
                # 处理聊天消息
                chat_request = ChatRequest(**data.get("data", {}))
                
                # 获取用户输入
                user_input = chat_request.messages[-1].content if chat_request.messages else "你好"
                
                # 准备RAG上下文（如果需要）
                if chat_request.query:
                    retriever = MilvusRetriever(milvus_service=milvus_service)
                    docs = retriever.get_relevant_documents(chat_request.query)
                    context = "\n\n".join([doc.page_content for doc in docs])
                    # 将RAG上下文添加到用户输入中
                    if context:
                        user_input = f"查询: {chat_request.query}\n\n相关上下文: {context}\n\n用户问题: {user_input}"
                
                # 注意：音频输出由后端处理，无需修改用户输入
                
                # 发送开始响应标记
                await websocket.send_json({
                    "type": "response_start",
                    "session_id": session_id
                })
                
                # 收集完整的响应内容和仅用于TTS的最终答案内容
                full_response = ""
                tts_content = ""  # 仅用于TTS的内容（只包含最终答案）
                final_answer_started = False
                
                # 流式发送响应
                async for chunk in tool_agent.astream_run(user_input, session_id=session_id):
                    chunk_content = chunk.get('content', '')
                    full_response += chunk_content
                    
                    # 检查是否是最终答案内容（用于TTS）
                    chunk_type = chunk.get('type', 'llm_stream')
                    if chunk_type == 'llm_stream':
                        # 这是经过agent.py处理的最终答案内容，用于TTS
                        tts_content += chunk_content
                    
                    await websocket.send_json({
                        "type": "response_chunk",
                        "content": chunk_content,
                        "content_type": "text",
                        "session_id": session_id
                    })
                
                # 发送结束标记
                await websocket.send_json({
                    "type": "response_end",
                    "finish_reason": "stop",
                    "session_id": session_id
                })
                
                # 如果需要音频输出，异步生成并发送语音（不阻塞UI）
                # 只对最终答案内容生成语音，不包括思考过程
                if chat_request.audio_output and tts_content.strip():
                    import asyncio
                    
                    async def generate_tts():
                        try:
                            # 检查WebSocket连接状态
                            if websocket.client_state.name != "CONNECTED":
                                return
                                
                            # 发送TTS开始状态
                            await websocket.send_json({
                                "type": "tts_start",
                                "message": "正在生成语音...",
                                "session_id": session_id
                            })
                            
                            # 限制文本长度以减小音频大小（最多500字符）
                            # 使用tts_content而不是full_response，只包含最终答案
                            tts_text = tts_content[:500] if len(tts_content) > 500 else tts_content
                            
                            # 生成语音文件（使用默认音色）
                            audio_url = await tts_service.text_to_speech_file(
                                text=tts_text
                            )
                            
                            # 再次检查WebSocket连接状态
                            if websocket.client_state.name != "CONNECTED":
                                return
                                
                            # 发送TTS完成状态，包含音频URL
                            await websocket.send_json({
                                "type": "tts_end",
                                "message": "语音生成完成",
                                "audio_url": audio_url,
                                "session_id": session_id
                            })
                            
                            # 注意：这里不立即标记为已播放，而是等待客户端播放完成后通知
                            
                        except Exception as tts_error:
                            # 只在连接有效时发送错误消息
                            try:
                                if websocket.client_state.name == "CONNECTED":
                                    await websocket.send_json({
                                        "type": "tts_error",
                                        "message": f"语音生成失败: {str(tts_error)}",
                                        "session_id": session_id
                                    })
                            except:
                                # 如果发送错误消息也失败，则忽略
                                pass
                    
                    # 创建异步任务，不等待完成
                    asyncio.create_task(generate_tts())
                
            elif message_type == "ping":
                # 处理心跳
                await websocket.send_json({
                    "type": "pong",
                    "timestamp": data.get("timestamp")
                })
                
            elif message_type == "clear_session":
                # 清除会话记忆
                memory_manager.clear_memory(session_id)
                await websocket.send_json({
                    "type": "session_cleared",
                    "session_id": session_id
                })
                
            elif message_type == "audio_played":
                # 处理音频播放完成通知
                audio_url = data.get("audio_url")
                if audio_url:
                    audio_manager.mark_as_played_by_url(audio_url)
                    await websocket.send_json({
                        "type": "audio_marked",
                        "message": "音频文件已标记为已播放",
                        "audio_url": audio_url,
                        "session_id": session_id
                    })
                else:
                    await websocket.send_json({
                        "type": "error",
                        "message": "缺少audio_url参数",
                        "session_id": session_id
                    })
                
    except WebSocketDisconnect:
        print(f"WebSocket连接断开: session_id={session_id}")
    except Exception as e:
        print(f"WebSocket错误: {e}")
        await websocket.send_json({
            "type": "error",
            "message": str(e),
            "session_id": session_id
        })
        await websocket.close()